import os
import shutil
import os.path as osp

root_dir = './helmat/'

if not osp.exists(osp.join(root_dir, 'annotations')):
    os.mkdir(osp.join(root_dir, 'annotations'))

if not osp.exists(osp.join(root_dir, 'train2017')):
    os.mkdir(osp.join(root_dir, 'train2017'))

if not osp.exists(osp.join(root_dir, 'test2017')):
    os.mkdir(osp.join(root_dir, 'test2017'))

if not osp.exists(osp.join(root_dir, 'val2017')):
    os.mkdir(osp.join(root_dir, 'val2017'))

with open(osp.join(root_dir, 'VOC2007', 'ImageSets', 'Main', 'train.txt'), 'r') as f:
    notes = f.readlines()
notes = [x.strip() for x in notes]

for note in notes:
    shutil.copy(osp.join(root_dir, 'VOC2007', 'JPEGImages', note+'.jpg'), osp.join(root_dir, 'trainval2017', note+'.jpg'))

with open(osp.join(root_dir, 'VOC2007', 'ImageSets', 'Main', 'test.txt'), 'r') as f:
    notes = f.readlines()
notes = [x.strip() for x in notes]

for note in notes:
    shutil.copy(osp.join(root_dir, 'VOC2007', 'JPEGImages', note+'.jpg'), osp.join(root_dir, 'test2017', note+'.jpg'))

with open(osp.join(root_dir, 'VOC2007', 'ImageSets', 'Main', 'val.txt'), 'r') as f:
    notes = f.readlines()
notes = [x.strip() for x in notes]

for note in notes:
    shutil.copy(osp.join(root_dir, 'VOC2007', 'JPEGImages', note+'.jpg'), osp.join(root_dir, 'val2017', note+'.jpg'))
